
import dotenv
from langchain_community.chat_models import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel

# 通过 runable parallel 实现两条链并行
dotenv.load_dotenv()
# 定义两个prompt
joke_prompt = ChatPromptTemplate.from_template("请讲一个关于{subject}的冷笑话")
poem_prompt = ChatPromptTemplate.from_template("请写一首关于{subject}的诗")
# 创建大预言模型
llm = ChatOpenAI(model= "gpt-3.5-turbo-16k")
# 创建解析器

parser = StrOutputParser()
# 编排链
joke_chain = joke_prompt|llm|parser
poem_chain = poem_prompt|llm|parser

# 并行链

map_chain = RunnableParallel(joke = joke_chain,poem = poem_chain)
res = map_chain.invoke({"subject": "程序员"})
print(res)